#!/usr/bin/env python
# -*- coding: utf-8 -*-
"""
统一的数据预处理脚本
功能：
- 从 CodeSearchNet 原始数据生成 complex 子集（单函数、去类、无 self/cls、AST 校验）
- 从 MBPP 数据生成 simple 子集（单函数、英文模板指令）
- 划分 train/valid/test 并保存到 data/processed 目录

说明：
- 输出文件：
    data/processed/complex_train_text2code.jsonl
    data/processed/complex_valid_text2code.jsonl
    data/processed/complex_test_text2code.jsonl
    data/processed/simple_train_text2code.jsonl
    data/processed/simple_valid_text2code.jsonl
    data/processed/simple_test_text2code.jsonl
"""

import os
import json
import ast
import warnings
from sklearn.model_selection import train_test_split
from tqdm import tqdm

# 忽略数据集中代码字符串的转义警告
warnings.filterwarnings('ignore', category=SyntaxWarning)

# ===================== #
# 路径配置
# ===================== #
RAW_DIR = "data/raw"
OUT_DIR = "data/processed"
os.makedirs(OUT_DIR, exist_ok=True)

def read_jsonl(path):
    """读取JSONL文件"""
    data = []
    with open(path, 'r', encoding='utf-8') as f:
        for line in f:
            try:
                data.append(json.loads(line))
            except json.JSONDecodeError:
                continue
    return data

def sanitize_docstring(text):
    """清理 docstring 中的特殊字符，避免 SyntaxWarning
    
    将未转义的反斜杠转换为双反斜杠，避免在后续处理中产生警告
    """
    if not text:
        return text

    replacements = {
        '\\m': '\\\\m',
        '\\s': '\\\\s', 
        '\\d': '\\\\d',
        '\\w': '\\\\w',
        '\\S': '\\\\S',
        '\\D': '\\\\D',
        '\\W': '\\\\W',
        '\\_': '\\\\_',
        '\\*': '\\\\*',
        '\\-': '\\\\-',
        '\\+': '\\\\+',
        '\\(': '\\\\(',
        '\\)': '\\\\)',
        '\\[': '\\\\[',
        '\\]': '\\\\]',
        '\\{': '\\\\{',
        '\\}': '\\\\}',
        '\\.': '\\\\.',
        '\\?': '\\\\?',
        '\\|': '\\\\|',
        '\\^': '\\\\^',
        '\\$': '\\\\$',
        '\\<': '\\\\<',
        '\\>': '\\\\>',
        '\\=': '\\\\=',
        '\\!': '\\\\!',
        '\\%': '\\\\%',
        '\\&': '\\\\&',
        '\\#': '\\\\#',
        '\\@': '\\\\@',
        '\\/': '\\\\/',
        '\\:': '\\\\:',
        '\\;': '\\\\;',
        '\\,': '\\\\,',
        '\\ ': '\\\\ ',  # 空格
        '\\c': '\\\\c',
        '\\e': '\\\\e',
        '\\f': '\\\\f',
        '\\g': '\\\\g',
        '\\h': '\\\\h',
        '\\i': '\\\\i',
        '\\j': '\\\\j',
        '\\k': '\\\\k',
        '\\l': '\\\\l',
        '\\o': '\\\\o',
        '\\p': '\\\\p',
        '\\q': '\\\\q',
        '\\u': '\\\\u',
        '\\v': '\\\\v',
        '\\x': '\\\\x',
        '\\y': '\\\\y',
        '\\z': '\\\\z',
        '\\A': '\\\\A',
        '\\B': '\\\\B',
        '\\C': '\\\\C',
        '\\E': '\\\\E',
        '\\F': '\\\\F',
        '\\G': '\\\\G',
        '\\H': '\\\\H',
        '\\I': '\\\\I',
        '\\J': '\\\\J',
        '\\K': '\\\\K',
        '\\L': '\\\\L',
        '\\M': '\\\\M',
        '\\N': '\\\\N',
        '\\O': '\\\\O',
        '\\P': '\\\\P',
        '\\Q': '\\\\Q',
        '\\R': '\\\\R',
        '\\T': '\\\\T',
        '\\U': '\\\\U',
        '\\V': '\\\\V',
        '\\X': '\\\\X',
        '\\Y': '\\\\Y',
        '\\Z': '\\\\Z',
    }
    
    # 直接用正则替换所有单反斜杠（但保留 \n, \t, \r 等有效转义）
    import re
    # 保护有效的转义序列
    valid_escapes = ['\\n', '\\t', '\\r', '\\\\', "\\'", '\\"', '\\a', '\\b', '\\f', '\\v']

    result = text
    for old, new in replacements.items():
        if old in result:
            result = result.replace(old, new)
    
    return result

def build_complex_samples(data, allow_varargs: bool = False, allow_param_names_args_kwargs_obj: bool = False):
    """转换 CodeSearchNet 数据集（英文） -> 单函数样本列表。
    - 方法->函数转换
    - 仅函数过滤（无类、无 self/cls、无 dunder、无 *args/**kwargs）
    - 生成统一英文模板 instruction
    """
    processed = []

    # --- 内联一份保守的 方法->函数 转换与过滤 ---
    class SelfUsageInspector(ast.NodeVisitor):
        def __init__(self):
            self.attr_loads = set()
            self.attr_stores = set()
            self.self_method_calls = set()
        def visit_Attribute(self, node: ast.Attribute):
            if isinstance(node.value, ast.Name) and node.value.id == "self":
                if isinstance(node.ctx, ast.Load):
                    self.attr_loads.add(node.attr)
                elif isinstance(node.ctx, (ast.Store,)):
                    self.attr_stores.add(node.attr)
            self.generic_visit(node)
        def visit_Call(self, node: ast.Call):
            if isinstance(node.func, ast.Attribute):
                if isinstance(node.func.value, ast.Name) and node.func.value.id == "self":
                    self.self_method_calls.add(node.func.attr)
            self.generic_visit(node)

    def _strip_docstring(fn: ast.FunctionDef):
        if fn.body and isinstance(fn.body[0], ast.Expr) and isinstance(getattr(fn.body[0], 'value', None), ast.Constant) and isinstance(fn.body[0].value.value, str):
            fn.body = fn.body[1:]

    def _replace_self_attr_with_params(fn: ast.FunctionDef, attr_names):
        class Replacer(ast.NodeTransformer):
            def visit_Attribute(self, node: ast.Attribute):
                if isinstance(node.value, ast.Name) and node.value.id == "self" and node.attr in attr_names:
                    return ast.copy_location(ast.Name(id=node.attr, ctx=node.ctx), node)
                return self.generic_visit(node)
        Replacer().visit(fn)

    def _remove_self_param(fn: ast.FunctionDef):
        if fn.args.args and fn.args.args[0].arg == "self":
            fn.args.args = fn.args.args[1:]

    def transform_method_to_function(code: str):
        try:
            tree = ast.parse(code)
        except SyntaxError:
            return None
        for node in tree.body:
            if isinstance(node, ast.FunctionDef):
                # 顶层普通函数，直接返回
                if not (node.args.args and node.args.args[0].arg == "self"):
                    _strip_docstring(node)
                    mod = ast.Module(body=[node], type_ignores=[])
                    ast.fix_missing_locations(mod)
                    try:
                        return ast.unparse(mod)
                    except Exception:
                        return None
                # 顶层"方法风格"(以 self 为首参) — 视作方法进行转换
                insp = SelfUsageInspector()
                insp.visit(node)
                if insp.self_method_calls:
                    return None
                if insp.attr_stores:
                    return None
                fn = ast.FunctionDef(
                    name=node.name,
                    args=ast.arguments(
                        posonlyargs=[],
                        args=[*node.args.args],
                        kwonlyargs=list(node.args.kwonlyargs),
                        kw_defaults=list(node.args.kw_defaults),
                        defaults=list(node.args.defaults),
                        vararg=node.args.vararg,
                        kwarg=node.args.kwarg,
                    ),
                    body=[*node.body],
                    decorator_list=[],
                    returns=node.returns,
                    type_comment=None,
                )
                _remove_self_param(fn)
                for a in sorted(insp.attr_loads):
                    fn.args.args.append(ast.arg(arg=a))
                _replace_self_attr_with_params(fn, set(insp.attr_loads))
                _strip_docstring(fn)
                mod = ast.Module(body=[fn], type_ignores=[])
                ast.fix_missing_locations(mod)
                try:
                    return ast.unparse(mod)
                except Exception:
                    return None
        for node in tree.body:
            if isinstance(node, ast.ClassDef):
                for m in node.body:
                    if isinstance(m, ast.FunctionDef) and (m.args.args and m.args.args[0].arg == "self"):
                        insp = SelfUsageInspector()
                        insp.visit(m)
                        if insp.self_method_calls:
                            continue
                        if insp.attr_stores:
                            continue
                        fn = ast.FunctionDef(
                            name=m.name,
                            args=ast.arguments(
                                posonlyargs=[],
                                args=[*m.args.args],
                                kwonlyargs=list(m.args.kwonlyargs),
                                kw_defaults=list(m.args.kw_defaults),
                                defaults=list(m.args.defaults),
                                vararg=m.args.vararg,
                                kwarg=m.args.kwarg,
                            ),
                            body=[*m.body],
                            decorator_list=[],
                            returns=m.returns,
                            type_comment=None,
                        )
                        _remove_self_param(fn)
                        for a in sorted(insp.attr_loads):
                            fn.args.args.append(ast.arg(arg=a))
                        _replace_self_attr_with_params(fn, set(insp.attr_loads))
                        _strip_docstring(fn)
                        mod = ast.Module(body=[fn], type_ignores=[])
                        ast.fix_missing_locations(mod)
                        try:
                            return ast.unparse(mod)
                        except Exception:
                            continue
        return None

    def is_single_function_strict(code: str) -> bool:
        s = (code or "").strip()
        if not s.startswith("def "):
            return False
        if "class " in s:
            return False
        try:
            tree = ast.parse(s)
        except SyntaxError:
            return False
        fdefs = [n for n in tree.body if isinstance(n, ast.FunctionDef)]
        if len(fdefs) != 1:
            return False
        fn = fdefs[0]
        if fn.name.startswith("__"):
            return False
        if not allow_varargs:
            if getattr(fn.args, 'vararg', None) is not None or getattr(fn.args, 'kwarg', None) is not None:
                return False
        bad_fn_names = {"get_python_obj", "get_python_args", "add_and_find_max", "add_max"}
        if fn.name in bad_fn_names:
            return False
        bad_params = {"self", "cls"}
        if not allow_param_names_args_kwargs_obj:
            bad_params.update({"obj", "args", "kwargs"})
        for a in fn.args.args:
            if a.arg in bad_params:
                return False
        return True

    def build_instruction(desc: str) -> str:
        desc = (desc or '').strip()
        return (
            f"Write a Python function that {desc}.\n"
            f"Constraints: Output only one complete Python function; no classes or methods; no 'self' or 'cls'; no dunder methods (e.g., __init__, __call__); no extra explanations or comments."
        )
    
    print(f"[INFO] 处理全部数据: {len(data)} 个样本...")
    
    for item in tqdm(data, desc="转换数据"):
        code = item.get("code", "").strip()
        doc_en = item.get("docstring", "").strip()
        
        # 清理 docstring 中的特殊字符
        doc_en = sanitize_docstring(doc_en)
        
        if not code or not doc_en:
            continue
        
        code_fn = code
        tr = transform_method_to_function(code)
        if tr is not None:
            code_fn = tr
        if not is_single_function_strict(code_fn):
            continue
        # 构造统一英文模板的 instruction（更贴近训练/推理分布）
        instruction = build_instruction(doc_en.strip())
        processed.append({
            "instruction": instruction,
            "output": code_fn,
            "task_type": "complex"
        })

    return processed

# ==== 共享工具（MBPP/Complex 共用） ====
def is_single_function_strict(code: str) -> bool:
    s = (code or '').strip()
    if not s.startswith('def '):
        return False
    if 'class ' in s:
        return False
    try:
        tree = ast.parse(s)
    except SyntaxError:
        return False
    fdefs = [n for n in tree.body if isinstance(n, ast.FunctionDef)]
    if len(fdefs) != 1:
        return False
    fn = fdefs[0]
    if fn.name.startswith('__'):
        return False
    if getattr(fn.args, 'vararg', None) is not None or getattr(fn.args, 'kwarg', None) is not None:
        return False
    bad_params = {'self', 'cls', 'obj', 'args', 'kwargs'}
    for a in fn.args.args:
        if a.arg in bad_params:
            return False
    return True

def build_instruction(desc: str) -> str:
    desc = (desc or '').strip().rstrip('.')
    return (
        f"Write a Python function that {desc}.\n"
        f"Constraints: Output only one complete Python function; no classes or methods; no 'self' or 'cls'; no dunder methods (e.g., __init__, __call__); no extra explanations or comments."
    )

def build_datasets(
    mbpp_path: str | None = None,
    salvage_preamble: bool = True,
    complex_allow_varargs: bool = True,
    complex_allow_param_names_args_kwargs_obj: bool = True,
):
    """处理并分离为 complex（大规模生产代码）和 simple（MBPP/算法题）子集。
    输出：
      data/processed/complex_train_text2code.jsonl
      data/processed/complex_valid_text2code.jsonl
      data/processed/complex_test_text2code.jsonl
      data/processed/simple_train_text2code.jsonl
      data/processed/simple_valid_text2code.jsonl
      data/processed/simple_test_text2code.jsonl
    """
    print("[INFO] 开始分离处理 complex/simple 数据子集...")

    # 1) 读取并处理复杂代码数据
    all_data = []
    files = [f for f in os.listdir(RAW_DIR) if f.endswith(".jsonl") and "python" in f]
    print(f"[INFO] 找到 {len(files)} 个数据文件用于 complex 子集")
    for file in files:
        path = os.path.join(RAW_DIR, file)
        file_data = read_jsonl(path)
        all_data.extend(file_data)
        print(f"[INFO] {file}: {len(file_data)} 个样本")

    complex_processed = build_complex_samples(
        all_data,
        allow_varargs=complex_allow_varargs,
        allow_param_names_args_kwargs_obj=complex_allow_param_names_args_kwargs_obj,
    )

    print(f"[INFO] complex 子集处理后样本数: {len(complex_processed)}")

    # 2) 读取并处理 MBPP/simple 数据
    simple_processed = []
    # 若未显式传入，则在常见位置自动探测
    if not mbpp_path:
        candidates = [
            os.path.join('data', 'raw', 'mbpp.jsonl'),
        ]
        mbpp_path = next((p for p in candidates if os.path.exists(p)), None)

    if mbpp_path and os.path.exists(mbpp_path):
        print(f"[INFO] 发现 MBPP 数据: {mbpp_path}，开始加载并处理为 simple 子集")
        mbpp = read_jsonl(mbpp_path)

        def _salvage_preamble_into_function(code: str) -> str | None:
            """若顶层仅包含 import/assign/expr-docstring + 单个函数，则将这些前导语句内联进函数体。
            返回新的“单函数字符串”，无法修复则返回 None。
            不改变函数签名，保持严格单函数输出格式。
            """
            s = (code or '').strip()
            if not s:
                return None
            try:
                tree = ast.parse(s)
            except SyntaxError:
                return None
            # 找到顶层函数；若多个函数或存在类，放弃修复
            top_defs = [n for n in tree.body if isinstance(n, ast.FunctionDef)]
            if len(top_defs) != 1:
                return None
            if any(isinstance(n, ast.ClassDef) for n in tree.body):
                return None
            fn = top_defs[0]
            # 收集可安全前置到函数体内的节点
            pre_nodes = []
            for n in tree.body:
                if n is fn:
                    continue
                if isinstance(n, (ast.Import, ast.ImportFrom, ast.Assign)):
                    pre_nodes.append(n)
                elif isinstance(n, ast.AnnAssign) and n.simple == 1 and isinstance(n.target, ast.Name):
                    pre_nodes.append(n)
                elif isinstance(n, ast.Expr) and isinstance(getattr(n, 'value', None), ast.Constant) and isinstance(n.value.value, str):
                    # 模块层 docstring，丢弃即可
                    continue
                else:
                    # 出现不可安全移动的节点，放弃修复
                    return None
            # 构造新的函数：在函数体最前面插入 pre_nodes 的深拷贝
            new_fn = ast.FunctionDef(
                name=fn.name,
                args=fn.args,
                body=[*pre_nodes, *fn.body],
                decorator_list=fn.decorator_list,
                returns=fn.returns,
                type_comment=getattr(fn, 'type_comment', None),
            )
            mod = ast.Module(body=[new_fn], type_ignores=[])
            ast.fix_missing_locations(mod)
            try:
                return ast.unparse(mod)
            except Exception:
                return None
        # MBPP 格式：尽量兼容 text/prompt/input/description 与 code/solution/reference_solution 等字段
        for item in tqdm(mbpp, desc="转换 MBPP"):
            # 尝试兼容多种字段名（mbpp 常见字段：text/code, prompt/completion, input/target）
            prompt = (
                item.get('instruction') or item.get('prompt') or item.get('input') or item.get('text') or item.get('description') or ''
            )
            output = (
                item.get('output') or item.get('completion') or item.get('target') or item.get('code') or item.get('solution') or item.get('reference_solution') or item.get('answer') or ''
            )
            if not prompt or not output:
                continue
            code = (output or '').strip()
            # 检查是否符合单函数规则
            if not is_single_function_strict(code):
                # 尝试修复前导 import/常量到函数体
                if salvage_preamble:
                    salvaged = _salvage_preamble_into_function(code)
                    if salvaged and is_single_function_strict(salvaged):
                        code = salvaged
                    else:
                        continue
                else:
                    continue
            # 统一指令模板
            instruction = prompt.strip()
            if not instruction.lower().startswith('write a python function'):
                instruction = build_instruction(instruction)
            simple_processed.append({
                "instruction": instruction,
                "output": code,
                "task_type": "simple"
            })
        print(f"[INFO] simple 子集（MBPP）处理后样本数: {len(simple_processed)}")
    else:
        print(f"[WARN] 未找到 MBPP 数据文件: {mbpp_path}，simple 子集将为空")

    # 3) 划分并保存两个子集
    def _save_splits(name, data):
        if not data:
            print(f"[WARN] 子集 {name} 数据为空，跳过保存")
            return
        train, temp = train_test_split(data, test_size=0.2, random_state=42)
        valid, test = train_test_split(temp, test_size=0.5, random_state=42)
        datasets = [("train", train), ("valid", valid), ("test", test)]
        for pname, dataset in datasets:
            out_path = os.path.join(OUT_DIR, f"{name}_{pname}_text2code.jsonl")
            with open(out_path, "w", encoding="utf-8") as f:
                for sample in dataset:
                    f.write(json.dumps(sample, ensure_ascii=False) + "\n")
            print(f"[SAVED] {out_path} - {len(dataset)} 个样本")

    _save_splits('complex', complex_processed)
    _save_splits('simple', simple_processed)

    print("[SUCCESS] complex/simple 子集预处理完成！")


if __name__ == "__main__":
    import argparse
    parser = argparse.ArgumentParser(description="生成 complex/simple 两个子集")
    parser.add_argument("--mbpp", dest="mbpp_path", default=None, help="MBPP 文件路径（留空将自动在 data/mbpp/ 与 data/raw/ 中探测）")
    args = parser.parse_args()
    
    # - MBPP: 自动内联前导 import/常量到函数体
    # - Complex: 允许 *args/**kwargs 与 args/kwargs/obj 参数名（仍禁止 self/cls）
    build_datasets(
        mbpp_path=args.mbpp_path,
        salvage_preamble=True,
        complex_allow_varargs=True,
        complex_allow_param_names_args_kwargs_obj=True,
    )
